package com.detector.overflow;

import com.detector.overflow.config.DetectorConfig;
import org.objectweb.asm.*;

/**
 * 方法访问器 - 监控方法中的数值运算
 *
 * @author detector
 */
public class OverflowMethodVisitor extends MethodVisitor {
    
    private String className;
    private String methodName;
    private String methodDescriptor;  // 方法描述符，用于区分重载方法
    private int currentLineNumber = -1;  // 当前行号
    private boolean isWhitelisted = false;  // 是否在白名单中
    
    public OverflowMethodVisitor(int api, MethodVisitor mv, String className, String methodName, String methodDescriptor) {
        super(api, mv);
        this.className = className;
        this.methodName = methodName;
        this.methodDescriptor = methodDescriptor;
        
        // 检查是否在白名单中
        String dotClassName = className.replace('/', '.');
        this.isWhitelisted = DetectorConfig.getInstance().isWhitelisted(dotClassName, methodName, this.methodDescriptor);
    }
    
    @Override
    public void visitLineNumber(int line, Label start) {
        // 记录当前行号
        this.currentLineNumber = line;
        super.visitLineNumber(line, start);
    }
    
    /**
     * 获取位置信息（包含行号）
     */
    private String getLocation() {
        String baseLocation = className.replace('/', '.') + "." + methodName;
        if (currentLineNumber > 0) {
            return baseLocation + ":" + currentLineNumber;
        }
        return baseLocation;
    }
    
    /**
     * 访问字节码指令，在遇到特定数值运算指令时插入溢出检测代码
     * 
     * 该方法会在ASM访问到方法中的每条指令时被调用，对于数值运算相关的指令，
     * 会在指令执行前插入相应的溢出检测调用。
     * 
     * @param opcode 指令的操作码，如 IADD(96), ISUB(100), IMUL(104) 等
     */
    @Override
    public void visitInsn(int opcode) {
        // 如果在白名单中，跳过所有检测
        if (isWhitelisted) {
            super.visitInsn(opcode);
            return;
        }
        
        if (opcode == Opcodes.IMUL) {
            // 在乘法前插入检查代码
            // 栈顶是两个int值，我们需要复制它们用于检查
            mv.visitInsn(Opcodes.DUP2); // 复制栈顶两个值
            
            // 调用检查方法: OverflowMonitor.checkMultiply(int a, int b, String location)
            mv.visitLdcInsn(getLocation());
            mv.visitMethodInsn(
                Opcodes.INVOKESTATIC,
                "com/detector/overflow/OverflowMonitor",
                "checkMultiply",
                "(IILjava/lang/String;)V",
                false
            );
        }
        // 监控整数加法指令 IADD
        else if (opcode == Opcodes.IADD) {
            mv.visitInsn(Opcodes.DUP2);
            mv.visitLdcInsn(getLocation());
            mv.visitMethodInsn(
                Opcodes.INVOKESTATIC,
                "com/detector/overflow/OverflowMonitor",
                "checkAdd",
                "(IILjava/lang/String;)V",
                false
            );
        }
        // 监控long乘法指令 LMUL
        else if (opcode == Opcodes.LMUL) {
            mv.visitInsn(Opcodes.DUP2_X2);  // 复制栈顶两个long值（每个long占2个栏位）
            mv.visitInsn(Opcodes.DUP2_X2);
            mv.visitLdcInsn(getLocation());
            mv.visitMethodInsn(
                Opcodes.INVOKESTATIC,
                "com/detector/overflow/OverflowMonitor",
                "checkMultiplyLong",
                "(JJLjava/lang/String;)V",
                false
            );
        }
        // 监控long加法指令 LADD
        else if (opcode == Opcodes.LADD) {
            mv.visitInsn(Opcodes.DUP2_X2);
            mv.visitInsn(Opcodes.DUP2_X2);
            mv.visitLdcInsn(getLocation());
            mv.visitMethodInsn(
                Opcodes.INVOKESTATIC,
                "com/detector/overflow/OverflowMonitor",
                "checkAddLong",
                "(JJLjava/lang/String;)V",
                false
            );
        }
        // 监控float乘法指令 FMUL
        else if (opcode == Opcodes.FMUL) {
            mv.visitInsn(Opcodes.DUP2);
            mv.visitLdcInsn(getLocation());
            mv.visitMethodInsn(
                Opcodes.INVOKESTATIC,
                "com/detector/overflow/OverflowMonitor",
                "checkMultiplyFloat",
                "(FFLjava/lang/String;)V",
                false
            );
        }
        // 监控float加法指令 FADD
        else if (opcode == Opcodes.FADD) {
            mv.visitInsn(Opcodes.DUP2);
            mv.visitLdcInsn(getLocation());
            mv.visitMethodInsn(
                Opcodes.INVOKESTATIC,
                "com/detector/overflow/OverflowMonitor",
                "checkAddFloat",
                "(FFLjava/lang/String;)V",
                false
            );
        }
        // 监控double乘法指令 DMUL
        else if (opcode == Opcodes.DMUL) {
            mv.visitInsn(Opcodes.DUP2_X2);
            mv.visitInsn(Opcodes.DUP2_X2);
            mv.visitLdcInsn(getLocation());
            mv.visitMethodInsn(
                Opcodes.INVOKESTATIC,
                "com/detector/overflow/OverflowMonitor",
                "checkMultiplyDouble",
                "(DDLjava/lang/String;)V",
                false
            );
        }
        // 监控double加法指令 DADD
        else if (opcode == Opcodes.DADD) {
            mv.visitInsn(Opcodes.DUP2_X2);
            mv.visitInsn(Opcodes.DUP2_X2);
            mv.visitLdcInsn(getLocation());
            mv.visitMethodInsn(
                Opcodes.INVOKESTATIC,
                "com/detector/overflow/OverflowMonitor",
                "checkAddDouble",
                "(DDLjava/lang/String;)V",
                false
            );
        }
        // 监控整数减法指令 ISUB
        else if (opcode == Opcodes.ISUB) {
            mv.visitInsn(Opcodes.DUP2);
            mv.visitLdcInsn(getLocation());
            mv.visitMethodInsn(
                Opcodes.INVOKESTATIC,
                "com/detector/overflow/OverflowMonitor",
                "checkSub",
                "(IILjava/lang/String;)V",
                false
            );
        }
        // 监控long减法指令 LSUB
        else if (opcode == Opcodes.LSUB) {
            mv.visitInsn(Opcodes.DUP2_X2);
            mv.visitInsn(Opcodes.DUP2_X2);
            mv.visitLdcInsn(getLocation());
            mv.visitMethodInsn(
                Opcodes.INVOKESTATIC,
                "com/detector/overflow/OverflowMonitor",
                "checkSubLong",
                "(JJLjava/lang/String;)V",
                false
            );
        }
        // 监控整数除法指令 IDIV
        else if (opcode == Opcodes.IDIV) {
            mv.visitInsn(Opcodes.DUP2);
            mv.visitLdcInsn(getLocation());
            mv.visitMethodInsn(
                Opcodes.INVOKESTATIC,
                "com/detector/overflow/OverflowMonitor",
                "checkDiv",
                "(IILjava/lang/String;)V",
                false
            );
        }
        // 监控long除法指令 LDIV
        else if (opcode == Opcodes.LDIV) {
            mv.visitInsn(Opcodes.DUP2_X2);
            mv.visitInsn(Opcodes.DUP2_X2);
            mv.visitLdcInsn(getLocation());
            mv.visitMethodInsn(
                Opcodes.INVOKESTATIC,
                "com/detector/overflow/OverflowMonitor",
                "checkDivLong",
                "(JJLjava/lang/String;)V",
                false
            );
        }
        // 监控类型转换指令 L2I (long → int)
        else if (opcode == Opcodes.L2I) {
            mv.visitInsn(Opcodes.DUP2);  // 复制long值（占用两个栏位）
            mv.visitLdcInsn(getLocation());
            mv.visitMethodInsn(
                Opcodes.INVOKESTATIC,
                "com/detector/overflow/OverflowMonitor",
                "checkL2I",
                "(JLjava/lang/String;)V",
                false
            );
        }
        // 监控类型转换指令 I2S (int → short)
        else if (opcode == Opcodes.I2S) {
            mv.visitInsn(Opcodes.DUP);
            mv.visitLdcInsn(getLocation());
            mv.visitMethodInsn(
                Opcodes.INVOKESTATIC,
                "com/detector/overflow/OverflowMonitor",
                "checkI2S",
                "(ILjava/lang/String;)V",
                false
            );
        }
        // 监控类型转换指令 I2B (int → byte)
        else if (opcode == Opcodes.I2B) {
            mv.visitInsn(Opcodes.DUP);
            mv.visitLdcInsn(getLocation());
            mv.visitMethodInsn(
                Opcodes.INVOKESTATIC,
                "com/detector/overflow/OverflowMonitor",
                "checkI2B",
                "(ILjava/lang/String;)V",
                false
            );
        }
        // 监控类型转换指令 I2C (int → char)
        else if (opcode == Opcodes.I2C) {
            mv.visitInsn(Opcodes.DUP);
            mv.visitLdcInsn(getLocation());
            mv.visitMethodInsn(
                Opcodes.INVOKESTATIC,
                "com/detector/overflow/OverflowMonitor",
                "checkI2C",
                "(ILjava/lang/String;)V",
                false
            );
        }
        // 监控类型转换指令 D2F (double → float)
        else if (opcode == Opcodes.D2F) {
            mv.visitInsn(Opcodes.DUP2);
            mv.visitLdcInsn(getLocation());
            mv.visitMethodInsn(
                Opcodes.INVOKESTATIC,
                "com/detector/overflow/OverflowMonitor",
                "checkD2F",
                "(DLjava/lang/String;)V",
                false
            );
        }
        // 监控类型转换指令 D2I (double → int)
        else if (opcode == Opcodes.D2I) {
            mv.visitInsn(Opcodes.DUP2);
            mv.visitLdcInsn(getLocation());
            mv.visitMethodInsn(
                Opcodes.INVOKESTATIC,
                "com/detector/overflow/OverflowMonitor",
                "checkD2I",
                "(DLjava/lang/String;)V",
                false
            );
        }
        // 监控类型转换指令 F2I (float → int)
        else if (opcode == Opcodes.F2I) {
            mv.visitInsn(Opcodes.DUP);
            mv.visitLdcInsn(getLocation());
            mv.visitMethodInsn(
                Opcodes.INVOKESTATIC,
                "com/detector/overflow/OverflowMonitor",
                "checkF2I",
                "(FLjava/lang/String;)V",
                false
            );
        }
        
        super.visitInsn(opcode);
    }
}
